#include <arm_neon.h>  // ARM NEON 头文件
#include <cassert>     // 用于 assert
#include <queue>       // 用于 priority_queue

std::priority_queue<std::pair<float, uint32_t>> neiji_neon(
    float* b1, float* b2, size_t base_number, size_t vecdim, size_t k) 
{
    std::priority_queue<std::pair<float, uint32_t>> q;
    assert(vecdim % 16 == 0);  // 确保向量维度是16的倍数

    for (size_t i = 0; i < base_number; ++i) {
        const float* vec = b1 + i * vecdim; // 当前数据库向量
        const float* query_vec = b2;        // 查询向量

        float32x4_t sum0 = vdupq_n_f32(0.0f);
        float32x4_t sum1 = vdupq_n_f32(0.0f);
        float32x4_t sum2 = vdupq_n_f32(0.0f);
        float32x4_t sum3 = vdupq_n_f32(0.0f);

        // 每次处理16个元素（4个NEON寄存器×4次加载）
        for (size_t j = 0; j < vecdim; j += 16) {
            // 加载数据库向量和查询向量的数据
            float32x4_t a0 = vld1q_f32(vec + j);
            float32x4_t b0 = vld1q_f32(query_vec + j);
            
            float32x4_t a1 = vld1q_f32(vec + j + 4);
            float32x4_t b1 = vld1q_f32(query_vec + j + 4);
            
            float32x4_t a2 = vld1q_f32(vec + j + 8);
            float32x4_t b2 = vld1q_f32(query_vec + j + 8);
            
            float32x4_t a3 = vld1q_f32(vec + j + 12);
            float32x4_t b3 = vld1q_f32(query_vec + j + 12);

            // 乘积累加
            sum0 = vmlaq_f32(sum0, a0, b0);
            sum1 = vmlaq_f32(sum1, a1, b1);
            sum2 = vmlaq_f32(sum2, a2, b2);
            sum3 = vmlaq_f32(sum3, a3, b3);
        }

        // 合并累加器
        sum0 = vaddq_f32(sum0, sum1);
        sum2 = vaddq_f32(sum2, sum3);
        sum0 = vaddq_f32(sum0, sum2);

        // 水平求和
        float32x2_t sum_low  = vget_low_f32(sum0);
        float32x2_t sum_high = vget_high_f32(sum0);
        sum_low = vadd_f32(sum_low, sum_high);
        float dis = vget_lane_f32(vpadd_f32(sum_low, sum_low), 0);

        // 将内积转换为距离
        dis = 1 - dis;  

        // 维护最小堆保存top-k
        if (q.size() < k) {
            q.emplace(dis, i);
        } else if (dis < q.top().first) {
            q.pop();
            q.emplace(dis, i);
        }
    }
    return q;
}